from fastapi import FastAPI, Request
from pydantic import BaseModel, Field
from wordcloud import WordCloud
import jieba
import matplotlib.pyplot as plt
import io
import base64
import os # 导入 os 模块，用于检查文件是否存在

app = FastAPI(title="词云生成服务")

# 定义请求体的模型
class TextInput(BaseModel):
    text: str
    # 将词云参数设为可选，并提供默认值
    width: int = Field(1920, description="词云图片的宽度，单位像素")
    height: int = Field(1080, description="词云图片的高度，单位像素")
    background_color: str = Field("white", description="词云图片的背景颜色")
    max_words: int = Field(6000, description="词云中显示的最大词语数量")
    min_font_size: int = Field(2, description="词云中词语的最小字体大小")
    prefer_horizontal: float = Field(0.6, description="词语水平排列的偏好 (0-1)")

# 定义词云字体路径（请根据你的实际情况修改！）
# 推荐使用支持中文的字体，例如 simhei.ttf 或 Microsoft YaHei 等
FONT_PATH = "./assets/SourceHanSerifSC-SemiBold.otf"

@app.post("/wordcloud")
async def generate_wordcloud(input_data: TextInput):
    """
    接收文本和可选的词云参数，生成词云图片并以 Base64 编码返回。
    """
    if not input_data.text:
        return {"error": "请求体中 'text' 字段不能为空。"}

    # 检查字体文件是否存在
    if not os.path.exists(FONT_PATH):
        return {"error": f"字体文件未找到: {FONT_PATH}。请确保字体文件存在于正确路径。"}

    try:
        # segmented_text = " ".join(jieba.cut(input_data.text))

        # 创建 WordCloud 对象，使用请求体中的参数（如果有），否则使用默认值
        wordcloud = WordCloud(
            font_path=FONT_PATH,
            width=input_data.width,
            height=input_data.height,
            background_color=input_data.background_color,
            max_words=input_data.max_words,
            min_font_size=input_data.min_font_size,
            prefer_horizontal=input_data.prefer_horizontal
        ).generate(input_data.text)
        # ).generate(segmented_text)  # 使用分词后的文本生成词云


        # 将词云图绘制到内存中
        # 注意：这里的 figsize 应该与 wordcloud 的 width/height 保持比例，否则可能拉伸
        # plt.figure(figsize=(input_data.width / 100, input_data.height / 100)) # 转换像素到英寸
        plt.figure(figsize=(input_data.width / plt.rcParams['figure.dpi'], input_data.height / plt.rcParams['figure.dpi']))
        plt.imshow(wordcloud, interpolation='bilinear')
        plt.axis("off") # 不显示坐标轴
        plt.tight_layout(pad=0) # 紧凑布局

        # 将 matplotlib 图形保存到内存缓冲区
        buffer = io.BytesIO()
        plt.savefig(buffer, format="png")
        plt.close() # 关闭图形，释放内存

        # 获取 Base64 编码的图片字符串
        image_base64 = base64.b64encode(buffer.getvalue()).decode("utf-8")

        return {
            "message": "词云生成成功！",
            "image": f"data:image/png;base64,{image_base64}" # 返回带 MIME 类型的 Base64 字符串
        }

    except Exception as e:
        # 打印更详细的错误信息到控制台，便于调试
        import traceback
        traceback.print_exc()
        return {"error": f"生成词云失败: {str(e)}。请检查输入参数或字体文件。"}

if __name__ == "__main__":
    import uvicorn
    # 在 8701 端口启动服务
    uvicorn.run(app, host="0.0.0.0", port=8701)